from abc import ABC, abstractmethod
from typing import Any, AsyncGenerator, List, Mapping, Sequence

from autogen_agentchat.base import ChatAgent, Response, TaskResult
from autogen_core import CancellationToken, ComponentBase
from pydantic import BaseModel
from autogen_agentchat.messages import BaseAgentEvent, BaseChatMessage, TextMessage, ModelClientStreamingChunkEvent
from autogen_agentchat.state import BaseState
class BaseChatAgent(ChatAgent, ABC, ComponentBase[BaseModel]):
    """ 聊天代理的基类
    这个抽象类为 :class:ChatAgent 提供了基础实现。若要创建新的聊天代理，需继承该类并实现 :meth:on_messages、:meth:on_reset 和 :attr:produced_message_types。如果需要流式处理，还需实现 :meth:on_messages_stream 方法。
    代理被视为有状态的，会在多次调用 :meth:on_messages 或 :meth:on_messages_stream 方法之间保持状态。代理应将状态存储在自身实例中，并实现 :meth:on_reset 方法以将代理重置为初始化状态。
    .. note::
    调用方在每次调用 :meth:on_messages 或 :meth:on_messages_stream 方法时，只需向代理传递新消息，切勿每次调用都传递完整的对话历史。创建新代理时必须遵循这一设计原则。
    """

    component_type = "agent"

    def __init__(self, name: str, description: str) -> None:
        """初始化代理，提供名称和描述。"""
        self._name = name
        if self._name.isidentifier() is False:
            raise ValueError("你的代理名称必须是有效的Python标识符。")
        self._description = description

    @property
    def name(self) -> str:
        """ 代理的名称，用于团队唯一标识代理。它应该在团队中是唯一的。 """
        return self._name

    @property
    def description(self) -> str:
        """代理的描述，用于团队决定使用哪个代理。它应该描述代理的能力和如何与它交互。"""
        return self._description

    @property
    @abstractmethod
    def produced_message_types(self) -> Sequence[type[BaseChatMessage]]:
        """
        返回代理在 :attr:`Response.chat_message` 字段中产生的消息类型。它们必须是 :class:`BaseChatMessage` 类型。
        """
        ...

    @abstractmethod
    async def on_messages(self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken) -> Response:
        """
        处理传入的消息并返回一个响应。
        笔记::
            代理是有状态的，传递给这个方法的消息应该是自上次调用以来的新消息。代理应该在调用之间维护其状态。例如，如果代理需要记住之前的消息来响应当前消息，它应该将之前的消息存储在代理状态中。
        """
        ...

    async def on_messages_stream(
        self, messages: Sequence[BaseChatMessage], cancellation_token: CancellationToken
    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | Response, None]:
        """
        处理传入的消息并返回一个消息流。
        笔记::
            代理是有状态的，传递给这个方法的消息应该是自上次调用以来的新消息。代理应该在调用之间维护其状态。例如，如果代理需要记住之前的消息来响应当前消息，它应该将之前的消息存储在代理状态中。
        """
        response = await self.on_messages(messages, cancellation_token)
        for inner_message in response.inner_messages or []:
            yield inner_message
        yield response

    async def run(
        self,
        *,
        task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> TaskResult:
        """
        运行代理并返回结果。
        """
        if cancellation_token is None:
            cancellation_token = CancellationToken()
        input_messages: List[BaseChatMessage] = []
        output_messages: List[BaseAgentEvent | BaseChatMessage] = []
        if task is None:
            pass
        elif isinstance(task, str):
            text_msg = TextMessage(content=task, source="user")
            input_messages.append(text_msg)
            output_messages.append(text_msg)
        elif isinstance(task, BaseChatMessage):
            input_messages.append(task)
            output_messages.append(task)
        else:
            if not task:
                raise ValueError("Task list cannot be empty.")
            # Task is a sequence of messages.
            for msg in task:
                if isinstance(msg, BaseChatMessage):
                    input_messages.append(msg)
                    output_messages.append(msg)
                else:
                    raise ValueError(f"Invalid message type in sequence: {type(msg)}")
        response = await self.on_messages(input_messages, cancellation_token)
        if response.inner_messages is not None:
            output_messages += response.inner_messages
        output_messages.append(response.chat_message)
        return TaskResult(messages=output_messages)

    async def run_stream(
        self,
        *,
        task: str | BaseChatMessage | Sequence[BaseChatMessage] | None = None,
        cancellation_token: CancellationToken | None = None,
    ) -> AsyncGenerator[BaseAgentEvent | BaseChatMessage | TaskResult, None]:
        """使用给定任务运行代理，并返回消息流，且流中的最后一项为最终任务结果。"""
        if cancellation_token is None:
            cancellation_token = CancellationToken()
        input_messages: List[BaseChatMessage] = []
        output_messages: List[BaseAgentEvent | BaseChatMessage] = []
        if task is None:
            pass
        elif isinstance(task, str):
            text_msg = TextMessage(content=task, source="user")
            input_messages.append(text_msg)
            output_messages.append(text_msg)
            yield text_msg
        elif isinstance(task, BaseChatMessage):
            input_messages.append(task)
            output_messages.append(task)
            yield task
        else:
            if not task:
                raise ValueError("Task list cannot be empty.")
            for msg in task:
                if isinstance(msg, BaseChatMessage):
                    input_messages.append(msg)
                    output_messages.append(msg)
                    yield msg
                else:
                    raise ValueError(f"Invalid message type in sequence: {type(msg)}")
        async for message in self.on_messages_stream(input_messages, cancellation_token):
            if isinstance(message, Response):
                yield message.chat_message
                output_messages.append(message.chat_message)
                yield TaskResult(messages=output_messages)
            else:
                yield message
                if isinstance(message, ModelClientStreamingChunkEvent):
                    # Skip the model client streaming chunk events.
                    continue
                output_messages.append(message)

    @abstractmethod
    async def on_reset(self, cancellation_token: CancellationToken) -> None:
        """重置代理到初始化状态。"""
        ...

    async def on_pause(self, cancellation_token: CancellationToken) -> None:
        """
        调用方在调用 :meth:on_messages 或 :meth:on_messages_stream 方法时，只需向代理传递新消息，
        """
        pass

    async def on_resume(self, cancellation_token: CancellationToken) -> None:
        """
        调用方在调用 :meth:on_messages 或 :meth:on_messages_stream 方法时，只需向代理传递新消息，
        """
        pass

    async def save_state(self) -> Mapping[str, Any]:
        """
        导出代理的当前状态。
        """
        return BaseState().model_dump()

    async def load_state(self, state: Mapping[str, Any]) -> None:
        """
        从保存的状态中恢复代理。默认实现适用于无状态代理。
        """
        BaseState.model_validate(state)

    async def close(self) -> None:
        """
        释放代理持有的任何资源。默认实现在 :class:`BaseChatAgent` 类中是一个空操作。子类可以重写此方法以实现自定义关闭行为。
        """
        pass